from setting import backtest
from setting import data
import polars as pl
import pickle
from setting.crate import *

# 策略名称 - 用于生成测试报告
strategy_name = "Alice1"

# 合约池
base_data_symbol = ["CFFEX_IF00_300"]  # 基准合约列表(用于时间轴)
trade_data_symbol = ["CFFEX_IF00_300", "CFFEX_IC00_300", "CFFEX_IH00_300"]  # 待交易合约,注意要交易时间要内包含于base,并且周期要保持和base一致
extra_data_symbol = ["CFFEX_IF00_86400", "CFFEX_IC00_86400", "CFFEX_IH00_86400"]  # 需增加的额外数据,用于跨周期调用

# 加载数据
base_data, all_trade_data = data.get_data(
    base_data_symbol,
    trade_data_symbol + extra_data_symbol,
    "20250101",
    "20251031"
)
# 分割数据
trade_data = {k: v for k, v in all_trade_data.items() if k in trade_data_symbol}
extra_data = {k: v for k, v in all_trade_data.items() if k in extra_data_symbol}

# 回测参数配置
initial_cash = 10000 * 10000.0  # 初始资金:1000万
fee_rate = 3.0 / 10000  # 手续费率:万分之三
# 读取费率
with open("setting/volume_multiple.pkl", "rb") as f:
    loaded_volume = pickle.load(f)
with open("setting/margin_ratio.pkl", "rb") as f:
    loaded_margin = pickle.load(f)
deposit_rates = loaded_margin
multipliers = loaded_volume

# 提前向量化指标
# 交易数据计算
for symbol, data in trade_data.items():
    trade_data[symbol] = data.lazy().with_columns([
        pl.col("close").rolling_mean(5).alias("ma5"),
        pl.col("close").rolling_mean(10).alias("ma10"),
        pl.col("close").rolling_mean(5).shift(1).alias("prev_ma5"),
        pl.col("close").rolling_mean(10).shift(1).alias("prev_ma10"),
    ]).collect()  # 一次性执行所有操作

# 处理额外数据
for symbol, data in extra_data.items():
    extra_data[symbol] = data.lazy().with_columns([
        pl.col("close").rolling_mean(5).alias("ma5"),
        pl.col("close").rolling_mean(10).alias("ma10"),
        pl.col("close").rolling_mean(5).shift(1).alias("prev_ma5"),
        pl.col("close").rolling_mean(10).shift(1).alias("prev_ma10"),
    ]).collect()

class Strategy(PythonStrategy):
    def __init__(self, backtest_instance,trade_data,extra_data):
        super().__init__(backtest_instance, trade_data, extra_data)

    def on_bar(self, timestamp):
        for symbol in self.trade_data.keys():
            recent_data = self.get_recent_data(symbol, timestamp, lookback=11, require_exact_match=True)
            if recent_data is not None:
                # self.cal_signal1(extract_symbol_base(symbol), recent_data, timestamp)
                self.calculate_signal(extract_symbol_base(symbol), recent_data,timestamp)

    def cal_signal1(self, symbol, recent_data, timestamp):
        """
        原始信号计算函数 - 动态计算技术指标
        
        Args:
            symbol: 品种代码
            recent_data: 最近的数据切片
            timestamp: 时间戳
        """
        # print(self.backtest_instance.today,self.backtest_instance.stock_profit,self.backtest_instance.profit)
        # print(self.backtest_instance.long_position,self.backtest_instance.short_position)
        # print("***********************************************************************")
        # 跨周期引用数据
        # df_daily = self.get_recent_data(symbol+"_86400",timestamp,2,False)
        # 动态计算移动平均线
        ma5 = recent_data["close"].rolling_mean(5)
        ma10 = recent_data["close"].rolling_mean(10)

        # 获取当前值(最后一个)
        current_ma5 = ma5[-1]
        current_ma10 = ma10[-1]
        # 获取前一个值(倒数第二个)
        prev_ma5 = ma5[-2]
        prev_ma10 = ma10[-2]
        # print(self.backtest_instance.today,current_ma5,current_ma10,prev_ma5,prev_ma10,recent_data["close"][-1])
        # 获取当前持仓
        current_holding = self.backtest_instance.long_position.get(symbol, 0)
        # 金叉信号:前一个周期MA5<=MA10,当前周期MA5>MA10,且无持仓时开多
        if prev_ma5 <= prev_ma10 and current_ma5 > current_ma10 and current_holding == 0:
            self.backtest_instance.buy(symbol, 7, recent_data["close"][-1])

        # 死叉信号:前一个周期MA5>=MA10,当前周期MA5<MA10,且有持仓时平仓
        elif prev_ma5 >= prev_ma10 and current_ma5 < current_ma10 and current_holding > 0:
            self.backtest_instance.sell(symbol, current_holding, recent_data["close"][-1])

    def calculate_signal(self, symbol, recent_data, timestamp):
        """
        优化信号计算函数 - 使用预计算的向量化指标
        
        Args:
            symbol: 品种代码
            recent_data: 最近的数据切片(包含预计算指标)
            timestamp: 时间戳
        """
        # 直接从预计算的数据中读取指标值(性能优化)
        current_ma5 = recent_data["ma5"][-1]  # 当前5周期均线
        current_ma10 = recent_data["ma10"][-1]  # 当前10周期均线
        prev_ma5 = recent_data["prev_ma5"][-1]  # 前一个5周期均线
        prev_ma10 = recent_data["prev_ma10"][-1]  # 前一个10周期均线
        # 获取当前持仓
        current_holding = self.backtest_instance.long_position.get(symbol, 0)

        # 金叉信号开多
        if prev_ma5 <= prev_ma10 and current_ma5 > current_ma10 and current_holding == 0:
            self.backtest_instance.buy(symbol, 2, recent_data["close"][-1])

        # 死叉信号平仓
        elif prev_ma5 >= prev_ma10 and current_ma5 < current_ma10 and current_holding > 0:
            self.backtest_instance.sell(symbol, current_holding, recent_data["close"][-1])


def main():
    """
    主函数 - 配置并运行回测
    """
    # 初始化回测引擎
    my_backtest = backtest.BackTest(initial_cash, fee_rate, deposit_rates, multipliers, trade_data)

    # 创建策略实例
    strategy = Strategy(my_backtest,trade_data,extra_data)

    # 准备回测时间序列
    datetime_strings = base_data.get_column("datetime").dt.strftime("%Y-%m-%d %H:%M:%S")

    # 启动回测
    # 参数说明:策略实例, 回测引擎, 时间字符串序列, 时间戳序列, 策略名称(用于生成报告)
    backtest.start_backtest(strategy, my_backtest, datetime_strings, base_data.get_column("timestamps"), strategy_name)


if __name__ == "__main__":
    main()
